""""""
from kuai.mol import Molecule

class MolecularVisitor(Exception):
    WHITE = 0
    GRAY = 1
    BLACK = 2
    
    def __init__(self):
        pass
    
    def start(self, mol, atom):
        pass
    
    def finish(self, mol, atom):
        pass
    
    def find(self, mol, atom, source):
        pass
    
    def back(self, mol, atom, source):
        pass
    
    
class SpanTreeVisitor(MolecularVisitor):
    def __init__(self):
        MolecularVisitor.__init__(self)
        self.__sources = {}
        
    def find(self, mol, atom, source):
        assert atom not in self.__sources
        self.__sources[atom] = source
        
        
    def get_source(self, atom):
        return self.__sources[atom]
    
    
def dft(mol, visitor, flags, start):
    visitor.start(mol, start)
    flags[start] = MolecularVisitor.GRAY
    for a, b in mol.neighbors[start]:
        if a in flags:
            visitor.back(mol, a, start)
        else:
            visitor.find(mol, a, start)
            dft(mol, visitor, flags, a)
    flags[start] = MolecularVisitor.BLACK
    visitor.finish(mol, start)
    
def bft(mol, visitor, flags, start):
    visitor.start(mol, start)
    queue = [start]
    while len(queue) > 0:
        start = queue.pop(0)
        visitor.start(mol, start)
        flags[start] = MolecularVisitor.GRAY
        for a, b in mol.neighbors[start]:
            if a in flags:
                visitor.back(mol, a, start)
            else:
                visitor.find(mol, a, start)
                queue.append(a)
        flags[start] = MolecularVisitor.BLACK
        visitor.finish(mol, start)


class FindConnectedPartVisitor(MolecularVisitor):
    """Used by split_molecule, do not use it directly"""
    def __init__(self):
        MolecularVisitor.__init__(self)
        self.__atoms = []
    
    def start(self, mol, atom):
        self.__atoms.append(atom)
    
    # def find(self, mol, atom, source):
    #     self.__atoms.append(atom)
        
    def get_result(self, mol):
        bonds = []
        for i in mol.bonds:
            if i.atom1 in self.__atoms and i.atom2 in self.__atoms:
                bonds.append(i)
                
        return Molecule(self.__atoms, bonds)
    
def split_molecule(mol):
    """It try to find all connected part of a molecule and return them in a list. 
Each connected part will be constructed as an independent molecule"""
    result = []
    flags = {}
    for i in mol.atoms:
        if i in flags:
            pass
        else:
            visitor = FindConnectedPartVisitor()
            dft(mol, visitor, flags, i)
            result.append(visitor.get_result(mol))
    return result

